#include <asm-generic/mman-common.h>
#include <sys/ioctl.h>
#include <acl/acl.h>
#include <rpg_dev.h>
#ifndef _GMEM_ASCEND_H
#define _GMEM_ASCEND_H

/* for hmadvise */
// #define SYS_hmadvise 441
// #define MADV_PREFETCH 32 /* prefetch pages for hNUMA node */
// TODO: Add further ifdef control to identify wether GMEM & hmadvise are supported

// #define MAP_PEER_SHARED 0X8000000

int gmemPrefetch(int fd, void *_device, size_t size, int hnid)
{
    struct hmadvise_arg hm_arg;
    hm_arg.start = (unsigned long)_device;
    hm_arg.len_in = size;
    hm_arg.hnid = hnid;
    hm_arg.behavior = MADV_PREFETCH;
    int ret;
    ret = ioctl(fd, GMEM_MADVISE, &hm_arg);
    if(ret != 0){
        perror("ioctl gmem_prefetch failed");
        printf("ioctl gmem_prefetch FAILED, ret:%d, gmem_fd:%d, start_addr:%p, size:%d, hnid:%d",
            ret,
            fd,
            _device,
            size,
            hnid
        )
        return ret;
    }
    printf("prefetch device buffer %lx, size %lu\n", (unsigned long)_device, size);
    return 0;
}

struct prefetch_buffer {
    int fd;
    void *device;
    size_t size;
    int hnid;
};

void _prefetch_helper(void *data)
{
    struct prefetch_buffer *buf = (struct prefetch_buffer *)data;
    printf("gmem_fd %d, prefetch addr %lx, size %lx, hnid %d\n",
        buf->fd,
        (unsigned long)buf->device,
        buf->size,
        buf->hnid
    );
    gmemPrefetch(buf->fd, buf->device, buf->size, buf->hnid);
    free(buf);
}

void gmemPrefetchAsync(int fd, void *device, size_t size, int hnid, aclrtStream stream)
{
    struct prefetch_buffer *buf = (struct prefetch_buffer *)malloc(sizeof(struct prefetch_buffer));
    buf->fd = fd;
    buf->device = device;
    buf->size = size;
    buf->hnid = hnid;
    aclrtLaunchCallback(_prefetch_helper, buf, ACL_CALLBACK_BLOCK, stream);
}

int get_gmem_fd(void)
{
    int gmem_fd = open("/dev/gmem", O_RDER);
    if(gmem_fd == -1){
        perror("Error opening file /dev/gmem");
        printf("opening file /dev/gmem FAILED\n");
        return -1;
    }
    return gmem_fd;
}

int get_hnid(int device_id)
{
    int id;
    struct gmem_hnid_arg arg;
    arg.hnuma_id = &id;

    int gmem_fd = get_gmem_fd();
    if(gmem_fd == -1){
        return -1;
    }

    if(ioctl(gmem_fd, GMEM_GET_HNIMA_ID, &arg) < 0){
        perror("ioctl GMEM_GET_HNIMA_ID");
        printf("ioctl GMEM_GET_HNIMA_ID FAILED\n");
        return -1;
    }
    return *arg.hnuma_id;
}

#endif